#include "header.h"



int init_network(in_addr_t server_ip) {
  struct sockaddr_in addr;
  int sockfd = socket(AF_INET, SOCK_STREAM, 0);
  if (sockfd == -1) return -1;

  memset(&addr, 0, sizeof(addr));
  addr.sin_addr.s_addr = server_ip;
  addr.sin_family = AF_INET;
  addr.sin_port = htons(SERVER_RECEIVE_PORT);

  if (connect(sockfd, &addr, sizeof(addr)) == -1) {
    close(sockfd);
    return -1;
  }
  return sockfd;
}



int init_data_stream(int sockfd) {
  int data_sockfd, addrlen;
  struct sockaddr_in addr;
  int data_port;

  data_sockfd = socket(AF_INET, SOCK_STREAM, 0);
  if (data_sockfd == -1) return -1;

  memset(&addr, 0, sizeof(addr));
  addr.sin_family = AF_INET;
  if (bind(data_sockfd, &addr, sizeof(addr)) == -1) {
    close(data_sockfd);
    return -1;
  }

  data_port = 0;
  while (data_port == 0) {
    if (getsockname(data_sockfd, &addr, &addrlen) == -1) {
      close(data_sockfd);
      return -1;
    }
    data_port = ntohs(addr.sin_port);
  }
  if (listen(data_sockfd, 2) == -1) {
    close(data_sockfd);
    return -1;
  }

  data_port = htonl(data_port);
  if (send_buf(sockfd, &data_port, 4) == -1) {
    close(data_sockfd);
    return -1;
  }

  data_sockfd = accept(data_sockfd, &addr, &addrlen);
  return data_sockfd;
}



int handle_get(int sockfd, int data_sockfd, char *remote, char *local) {
  char *buffer = malloc(BUFFER_SIZE);
  int len;
  FILE *file;
  char op;

  file = fopen(local, "wb");
  if (file == 0) {
    return -1;
  }
  op = 1;
  if (send_buf(sockfd, &op, 1) == -1) {
    fclose(file);
    return -1;
  }
  len = htonl(strlen(remote));
  if (send_buf(sockfd, &len, 4) == -1 ||
      send_buf(sockfd, remote, strlen(remote)) == -1 ||
      recv_buf(sockfd, &len, 4) <= 0) {
    fclose(file);
    return -1;
  }
  len = ntohl(len);
  if (len == -1) {
    fclose(file);
    return 0;
  }
  while (len) {
    if (len > BUFFER_SIZE) {
      if (recv_buf(data_sockfd, buffer, BUFFER_SIZE) <= 0) {
        fclose(file);
        return -1;
      }
      fwrite(buffer, 1, BUFFER_SIZE, file);
      len -= BUFFER_SIZE;
    } else {
      if (recv_buf(data_sockfd, buffer, len) <= 0) {
        fclose(file);
        return -1;
      }
      fwrite(buffer, 1, len, file);
      len = 0;
    }
  }
  fclose(file);
  return 1;
}



int handle_put(int sockfd, int data_sockfd, char *local, char *remote) {
  char *buffer = malloc(BUFFER_SIZE);
  int len, ack;
  FILE *file;
  char op;

  file = fopen(local, "rb");
  if (file == 0) {
    return 0;
  }
  op = 2;
  if (send_buf(sockfd, &op, 1) == -1) {
    fclose(file);
    return -1;
  }
  len = htonl(strlen(remote));
  if (send_buf(sockfd, &len, 4) == -1 ||
      send_buf(sockfd, remote, strlen(remote)) == -1 ||
      recv_buf(sockfd, &ack, 4) <= 0) {
    fclose(file);
    return -1;
  }
  ack = ntohl(ack);
  if (ack == -1) {
    fclose(file);
    return -1;
  }

  fseek(file, 0, SEEK_END);
  len = ftell(file);
  fseek(file, 0, SEEK_SET);
  len = htonl(len);
  if (send_buf(sockfd, &len, 4) == -1) {
    fclose(file);
    return -1;
  }
  len = ntohl(len);

  while (len) {
    if (len > BUFFER_SIZE) {
      fread(buffer, 1, BUFFER_SIZE, file);
      if (send_buf(data_sockfd, buffer, BUFFER_SIZE) == -1) {
        fclose(file);
        return -1;
      }
      len -= BUFFER_SIZE;
    } else {
      fread(buffer, 1, len, file);
      if (send_buf(data_sockfd, buffer, len) == -1) {
        fclose(file);
        return -1;
      }
      len = 0;
    }
  }
  fclose(file);
  return 1;
}


char* handle_pwd(int sockfd) {
  char op, *buf;
  int len;

  op = 3;
  if (send_buf(sockfd, &op, 1) == -1) {
    return 0;
  }
  if (recv_buf(sockfd, &len, 4) <= 0) {
    return 0;
  }
  len = ntohl(len);
  buf = malloc(len + 1);
  if (recv_buf(sockfd, buf, len) <= 0) {
    free(buf);
    return 0;
  }
  buf[len] = 0;
  return buf;
}


char* handle_dir(int sockfd) {
  char op, *buf;
  int len;

  op = 4;
  if (send_buf(sockfd, &op, 1) == -1) {
    return 0;
  }
  if (recv_buf(sockfd, &len, 4) <= 0) {
    return 0;
  }
  len = ntohl(len);
  buf = malloc(len + 1);
  if (recv_buf(sockfd, buf, len) <= 0) {
    free(buf);
    return 0;
  }
  buf[len] = 0;
  return buf;
}



char* handle_cd(int sockfd, char *path) {
  char op, *buf;
  int len = htonl(strlen(path));

  op = 5;
  if (send_buf(sockfd, &op, 1) == -1 ||
      send_buf(sockfd, &len, 4) == -1 ||
      send_buf(sockfd, path, ntohl(len)) == -1) {
    return 0;
  }
  if (recv_buf(sockfd, &len, 4) <= 0) {
    return 0;
  }
  len = ntohl(len);
  buf = malloc(len + 1);
  if (recv_buf(sockfd, buf, len) <= 0) {
    free(buf);
    return 0;
  }
  buf[len] = 0;
  return buf;
}



int main(int argc, char* argv[]) {
  in_addr_t server_ip = 0;
  int sockfd, data_sockfd, tmp, len;
  char *command, *s1, *s2, *s3, *str;

  if (argc == 2) server_ip = inet_addr(argv[1]);
  if (server_ip == 0) {
    printf("Usage: ./ftp_client.out [ip]\n");
    printf("Example: ./ftp_client.out 192.168.1.1\n");
    return 0;
  }

  sockfd = init_network(server_ip);
  if (sockfd == -1) {
    printf("Fail: init_network\n");
    return 0;
  }
  data_sockfd = init_data_stream(sockfd);
  if (data_sockfd == -1) {
    printf("Fail: init_data_stream\n");
    return 0;
  }

  command = malloc(BUFFER_SIZE);
  s1 = malloc(BUFFER_SIZE);
  s2 = malloc(BUFFER_SIZE);
  s3 = malloc(BUFFER_SIZE);
  while (1) {
    printf(">");
    fgets(command, BUFFER_SIZE - 1, stdin);
    len = sscanf(command, "%s%s%s", s1, s2, s3);
    if (strcmp(s1, "get") == 0 && len == 3) {
      tmp = handle_get(sockfd, data_sockfd, s2, s3);
      if (tmp == -1) {
        printf("Fail: get\n");
        break;
      } else if (tmp == 0) {
        printf("Error: remote file not found\n");
      } else {
        printf("Success: get %s %s\n", s2, s3);
      }
    } else if (strcmp(s1, "put") == 0 && len == 3) {
      tmp = handle_put(sockfd, data_sockfd, s2, s3);
      if (tmp == -1) {
        printf("Fail: put\n");
        break;
      } else if (tmp == 0) {
        printf("Error: local file not found\n");
      } else {
        printf("Success: put %s %s\n", s2, s3);
      }
    } else if (strcmp(s1, "pwd") == 0 && len == 1) {
      str = handle_pwd(sockfd);
      if (str == 0) {
        printf("Fail: pwd\n");
        break;
      } else {
        printf("%s\n", str);
        free(str);
      }
    } else if (strcmp(s1, "dir") == 0 && len == 1) {
      str = handle_dir(sockfd);
      if (str == 0) {
        printf("Fail: dir\n");
        break;
      } else {
        printf("%s\n", str);
        free(str);
      }
    } else if (strcmp(s1, "cd") == 0 && len == 2) {
      str = handle_cd(sockfd, s2);
      if (str == 0) {
        printf("Fail: cd\n");
        break;
      } else {
        printf("Current path: %s\n", str);
        free(str);
      }
    } else if (strcmp(s1, "?") == 0 && len == 1) {
      printf("Help:\n");
      printf("\tget [remote filename] [local filename]\n");
      printf("\tput [local filename] [remote filename]\n");
      printf("\tpwd\n");
      printf("\tdir\n");
      printf("\tcd [path]\n");
      printf("\t?\n");
      printf("\tquit\n");
    } else if (strcmp(s1, "quit") == 0 && len == 1) {
      break;
    } else {
      printf("Error: unkown command\n");
    }
  }

  return 0;
}
